﻿/*
* ////////////////////////////////////////////////
*
* 开发者：AZPEP
* 友情链接：
*	-	B站：		零分的考卷			https://space.bilibili.com/257391258
*	-	Github：	AZPEP-BiliBili		https://github.com/AZPEP-BiliBili
*	-	Gitee：		AZPEP_BiliBili		https://gitee.com/AZPEP_BiliBili
*
* CopyRight AZPEP
*
* ////////////////////////////////////////////////
*/

#include "pch.h"
#include "GetDrive.h"
#include "GetGPT.h"

/// <summary>
/// 返回计算机上所有GPT分区的信息
/// </summary>
/// <param name="PartitionInformation">指向所有GPT分区的信息的结构体指针的列表</param>
/// <returns>成功返回0，PartitionInformation指向所有GPT分区的信息的结构体指针的列表；失败返回GetLastError()的错误代码，PartitionInformation为空</returns>
DWORD MountEsp_GetAllGptPartitionInformationW(
	_Out_ std::vector<PartitionInformationW*>* PartitionInformation
) {
	DWORD dStatus;															//状态

	//获取DRIVE_LAYOUT_INFORMATION_EX
	std::vector<DRIVE_LAYOUT_INFORMATION_EX*> vec_DLIE;
	dStatus = MountEsp_GetAllDriveLayoutEx(&vec_DLIE);
	if (dStatus != 0) { PartitionInformation->clear(); return dStatus; }	//获取失败，返回

	//开始遍历处理_DLIE
	for (int i = 0; i < vec_DLIE.size(); i++) {
		DRIVE_LAYOUT_INFORMATION_EX* P_DLIE = vec_DLIE[i];
		std::vector<PartitionInformationW*> tmp_vec_PIW;					//存放一个驱动器的所有分区的信息
		dStatus = MountEsp_ParseAGptDriveLayoutInformationEx(
			i,
			P_DLIE,
			&tmp_vec_PIW
		);
		switch (dStatus) {
			//磁盘无分区/MBR磁盘，跳过处理
		case 20000:
		case 20001:
			continue;
			break;

			//处理成功，开始推入PartitionInformation
		case 0:
			for (int i = 0; i < tmp_vec_PIW.size(); i++) { PartitionInformation->push_back(tmp_vec_PIW[i]); }
			break;

			//获取失败，返回
		default:
			return dStatus;
		}
	}

	return 0;
}

/// <summary>
/// 处理一张GPT磁盘上所有分区的基本信息
/// </summary>
/// <param name="DLIE">指向DRIVE_LAYOUT_INFORMATION_EX的指针</param>
/// <param name="PIW">指向接受PartitionInformationW*的列表的指针</param>
/// <returns>成功返回0，同时PIW指向该硬盘上所有分区的基本信息；失败返回GetLastError()获得的错误代码，磁盘为MBR磁盘或未初始化返回20000，磁盘无分区返回20001；若返回值为非零值，则PIW为空</returns>
DWORD MountEsp_ParseAGptDriveLayoutInformationEx(
	_In_ int DiskNumber,
	_In_ DRIVE_LAYOUT_INFORMATION_EX* DLIE,
	_Out_ std::vector <PartitionInformationW*>* PIWL
) {
	//磁盘为MBR磁盘或未初始化返回20000
	if (DLIE->PartitionStyle == PARTITION_STYLE_MBR || DLIE->PartitionStyle == PARTITION_STYLE_RAW) return 20000;

	//磁盘无分区返回20001
	if (DLIE->PartitionCount == 0) return 20001;

	//遍历读取分区信息
	for (int i = 0; i < DLIE->PartitionCount; i++) {
		MallocEx(PartitionInformationW*, P_PIW, sizeof(PARTITION_INFORMATION_EX));

		//获取磁盘信息
		PARTITION_INFORMATION_EX* P_PIE = &DLIE->PartitionEntry[i];

		//获取分区信息
		DWORD dStatus = MountEsp_ParseAGptPartitionInformationEx(
			DiskNumber,
			P_PIE,
			P_PIW
		);

		//失败，返回
		if (dStatus != 0) { PIWL->clear(); return GetLastError(); }

		//否则，添加到PIWL
		PIWL->push_back(P_PIW);
	}

	//成功，返回
	return 0;
}

/// <summary>
/// 处理一个分区的基本信息
/// </summary>
/// <param name="DiskNumber">分区所在磁盘序号</param>
/// <param name="PIE">指向PARTITION_INFORMATION_EX的指针</param>
/// <param name="PIW">指向接受PartitionInformationW的指针，请事先使用malloc与ZeroMemory分配内存</param>
/// <returns>成功返回0，同时PIW指向PartitionInformationW的指针；失败返回GetLastError()获得的错误代码，同时PIW为空残缺的结构体</returns>
DWORD MountEsp_ParseAGptPartitionInformationEx(
	_In_ DWORD DiskNumber,
	_In_ PARTITION_INFORMATION_EX* PIE,
	_Out_ PartitionInformationW* PIW
) {
	//对PIL赋值
	PIW->DiskNumber = DiskNumber;					//磁盘序号
	PIW->PartitionNumber = PIE->PartitionNumber;	//分区序号
	PIW->PartitionId = PIE->Gpt.PartitionId;		//分区ID
	PIW->PartitionType = PIE->Gpt.PartitionType;	//分区类型
	PIW->PartitionStyle = PIE->PartitionStyle;		//分区样式

	//分区路径
	std::wstringstream wss_PartitionPath;
	wss_PartitionPath << L"\\Device\\Harddisk" << PIW->DiskNumber << L"\\Partition" << PIW->PartitionNumber;
	PIW->PartitionPath = wss_PartitionPath.str();

	//分区卷路径
	std::wstringstream wss_PartitionVolumePath;
	wss_PartitionVolumePath << L"\\\\?\\Volume" << GUID2StringW(PIW->PartitionId) << L"\\";
	PIW->PartitionVolumePath = wss_PartitionVolumePath.str();

	//获取盘符
	WCHAR Letter = GetVolumeMountLetterW(PIW->PartitionVolumePath);

	switch (Letter) {
		//失败
	case L'\1':
		return GetLastError();
		break;

		//没有盘符
	case L'\0':
		PIW->IsMounted = FALSE;
		PIW->Letter = NULL;
		break;

		//有盘符
	default:
		PIW->IsMounted = TRUE;
		PIW->Letter = Letter;
	}

	//处理驱动器类型

	return 0;
}

/// <summary>
/// 将GUID转换为std::wstring
/// </summary>
/// <param name="guid"></param>
wchar_t* GUID2StringW(GUID guid) {
	wchar_t* wc = (wchar_t*)malloc(MAX_PATH);

	wsprintf(wc, L"{%08x-%04x-%04x-%02x%02x-%02x%02x%02x%02x%02x%02x}", guid.Data1, guid.Data2, guid.Data3, guid.Data4[1], guid.Data4[0], guid.Data4[7], guid.Data4[6], guid.Data4[5], guid.Data4[4], guid.Data4[3], guid.Data4[2]);

	return wc;
}

/// <summary>
/// 通过卷GUID路径获取挂载的盘符。再次实现而非通过GetVolumePathNamesForVolumeName是因为GetVolumePathNamesForVolumeName存在诸多缺陷
/// </summary>
/// <param name="VolumePath">卷GUID路径</param>
/// <returns>成功返回盘符，卷未挂载返回\0，获取失败返回\1，可通过GetLastError获取错误信息</returns>
WCHAR GetVolumeMountLetterW(std::wstring VolumePath) {
	std::vector<std::wstring> MountedVolumeGuidPathList;	//存放已挂载盘符的GUID的列表

	std::vector<WCHAR> VolumeLetter = GetVolumeLetter(); //获取所有已挂载分区的盘符

	//通过遍历获取已挂载盘符的GUID，并存放在MountedVolumeGuidPath中
	for (int i = 0; i < VolumeLetter.size(); i++) {
		//拼接字符串，拼接后字符串：X:\ (X代表卷标)
		std::wstring LetterPath;
		LetterPath += VolumeLetter[i];
		LetterPath += L":\\";

		LPCWSTR LetterPath_LPW = LetterPath.c_str();
		//获取卷GUID
		LPWSTR VolumeGuidPath = (LPWSTR)malloc(MAX_PATH);
		ZeroMemory(VolumeGuidPath, MAX_PATH);
		BOOL Status = GetVolumeNameForVolumeMountPointW(
			LetterPath_LPW,
			VolumeGuidPath,
			MAX_PATH
		);

		//获取失败，报错ERROR_GETVOLUME_GET_VOLUME_MOUNT_PATH_FAIL并返回
		if (!Status) {
			free(VolumeGuidPath); //获取失败后释放掉占用的内存，防止爆栈。
			DWORD ErrorCode = GetLastError();
			SetLastError(ErrorCode);
			return L'\1';
		}

		//将VolumeGuidPath添加到MountedVolumeGuidPathList
		std::wstring MountedVolumeGuidPath = VolumeGuidPath;
		MountedVolumeGuidPathList.push_back(MountedVolumeGuidPath);
	}

	//遍历MountedVolumeGuidPathList，查找对应GUID
	for (int i = 0; i < MountedVolumeGuidPathList.size(); i++) {
		if (MountedVolumeGuidPathList[i].substr(0, 30) == VolumePath.substr(0, 30)) {
			return VolumeLetter[i];
		}
	}
	//跳出循环说明没找到，则返回0代表不存在
	return L'\0';
}

typedef struct SimplePartitionInfo {
	char Letter;					//盘符
	std::wstring VolumeGuidPath;	//卷GUID路径
};
/// <summary>
///  获取磁盘上所有已使用的盘符
/// </summary>
/// <returns>成功GetLastError()=0，失败GetLastError()为非零值</returns>
std::vector<WCHAR> GetVolumeLetter() {
	std::vector<WCHAR> DriveLetter;

	//创建完全盘符列表
	std::vector<WCHAR> AllLetter;
	for (WCHAR wc = L'A'; wc <= L'Z'; wc++) {
		std::wstringstream LetterPath;
		LetterPath << wc << L":/";
		/*
		//尝试打开磁盘
		HANDLE hDevice = CreateFile(
			LetterPath.str().c_str(),
			GENERIC_READ,
			FILE_SHARE_READ | FILE_SHARE_WRITE,
			NULL,
			OPEN_EXISTING,
			0,
			NULL
		);
		*/
		if (!(_waccess(LetterPath.str().c_str(), 0) == -1)) {
			//CloseHandle(hDevice);
			DriveLetter.push_back(wc);
		}
	}

	if (DriveLetter.size() == 0) {
		return (std::vector<WCHAR>)NULL;
	}

	SetLastError(0);
	return DriveLetter;
}